package com.ustc.hewei.snowflakemodule.snowFlake;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.ustc.hewei.snowflakemodule.snowFlake.entity.WorkIdNode;
import com.ustc.hewei.snowflakemodule.exception.TimestampException;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.RetryUntilElapsed;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.List;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;

/**
 * @author hewei
 * @version 1.0
 * @description: 使用zk生成唯一workId
 * @date 2022/10/12 19:41
 */

public class SnowflakeZkHandler {
    private static final Logger logger = LoggerFactory.getLogger(SnowflakeZkHandler.class);
    private final String zkAddress;
    private final String initPath = "/snowflake/workId";
    // 根据这个ip和port来分配机器码
    private String ip;
    private String port;
    private int workId;
    // 用ip和port拼接成的地址
    private final String listenedAddress;
    // 在zk中的节点全路径 initPath/listenedAddress-0000000/
    private String allAddress;
    // 上一次更新的时间
    private long lastUpdateTime;
    private ScheduledExecutorService executorService;

    public SnowflakeZkHandler(String zkAddress, String ip, String port) {
        this.ip = ip;
        this.zkAddress = zkAddress;
        this.port = port;
        listenedAddress = ip + ":" + port;
    }

    public boolean init() {
        try {
            CuratorFramework curator = createCurator(zkAddress, new RetryUntilElapsed(1000, 4), 10000, 6000);
            curator.start();
            // 判断根节点是否存在
            Stat stat = curator.checkExists().forPath(initPath);
            if (stat != null) {
                // 判断是否存在此地址的节点数据
                // 获取根节点下所有的子节点
                List<String> list = curator.getChildren().forPath(initPath);
                HashMap<String, Integer> workIdMap = new HashMap<>();
                HashMap<String, String> childPathMap = new HashMap<>();
                for (String childPath : list) {
                    String[] split = childPath.split("-");
                    workIdMap.put(split[0], Integer.parseInt(split[1]));
                    childPathMap.put(split[0], childPath);
                }
                Integer workIdTemp = workIdMap.get(listenedAddress);
                if (workIdTemp == null) {
                    // 不存在这个地址的子节点，直接创建新的
                    allAddress = createNode(curator);
                    String[] split = allAddress.split("-");
                    workId = Integer.parseInt(split[1]);
                    // 定时上报时间
                    scheduledUpdateNode(curator, allAddress);
                    logger.info("新建节点 {} 成功", allAddress);
                } else {
                    // 存在子节点，直接用上次分配的workId
                    workId = workIdTemp;
                    allAddress = initPath + "/" + childPathMap.get(listenedAddress);
                    // 检测上一次该地址的更新时间，判断是否时间回拨
                    if (!checkTimestamp(curator, allAddress)) {
                        throw new TimestampException("节点时间晚于当前时间!");
                    }
                    scheduledUpdateNode(curator, allAddress);
                    logger.info("更新旧节点 {} 成功", allAddress);
                }
            } else {
                // 根节点不存在，直接创建新节点
                allAddress = createNode(curator);
                // 在目标路径定时上报当前时间
                scheduledUpdateNode(curator, allAddress);
                logger.info("创建根节点 {} 成功", allAddress);
            }
        } catch (Exception e) {
            logger.error(e.getMessage());
            return false;
        }
        return true;
    }

    private boolean checkTimestamp(CuratorFramework curator, String path) throws Exception {
        byte[] bytes = curator.getData().forPath(path);
        WorkIdNode node = deSerializerNode(new String(bytes));
        return node.getTimeStamp() < System.currentTimeMillis();
    }

    private WorkIdNode deSerializerNode(String json) throws JsonProcessingException {
        ObjectMapper mapper = new ObjectMapper();
        return mapper.readValue(json, WorkIdNode.class);
    }

    private CuratorFramework createCurator(String zkAddress, RetryPolicy retryPolicy, int connectionTimeoutMs, int sessionTimeoutMs) {
        return CuratorFrameworkFactory.builder().connectString(zkAddress)
                .retryPolicy(retryPolicy)
                .sessionTimeoutMs(sessionTimeoutMs)
                .connectionTimeoutMs(connectionTimeoutMs).build();
    }

    private String createNode(CuratorFramework curator) throws Exception {
        try {
            // CreateMode.PERSISTENT_SEQUENTIAL设置序号自增模式，从零开始
            return curator.create().creatingParentsIfNeeded()
                    .withMode(CreateMode.PERSISTENT_SEQUENTIAL)
                    .forPath(initPath + "/" + listenedAddress + "-", buildData().getBytes());
        } catch (Exception e) {
            logger.error("创建节点失败 {}", e.getMessage());
            throw e;
        }
    }

    private String buildData() throws JsonProcessingException {
        WorkIdNode node = new WorkIdNode(ip, port, System.currentTimeMillis());
        ObjectMapper mapper = new ObjectMapper();
        return mapper.writeValueAsString(node);
    }

    private void scheduledUpdateNode(CuratorFramework curator, String path) {
        executorService = Executors.newSingleThreadScheduledExecutor(new ThreadFactory() {
            @Override
            public Thread newThread(Runnable r) {
                Thread thread = new Thread(r, "scheduled-update-node-" + port);
                // 设置守护线程
                thread.setDaemon(true);
                return thread;
            }
        });
        // 每3s更新一次节点时间
        executorService.scheduleWithFixedDelay(() -> {
            updateNodeTime(curator, path);
        }, 1L, 3L, TimeUnit.SECONDS);
    }

    public void updateNodeTime(CuratorFramework curator, String path) {
        try {
            if (System.currentTimeMillis() < lastUpdateTime) {
                return;
            }
            curator.setData().forPath(path, buildData().getBytes());
            lastUpdateTime = System.currentTimeMillis();
        } catch (Exception e) {
            logger.info("更新路径 {} 失败, 异常为 {}", path, e);
        }
    }

    public int getWorkId() {
        return workId;
    }
}
